package srt.tool;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import srt.ast.AssertStmt;
import srt.ast.AssignStmt;
import srt.ast.BinaryExpr;
import srt.ast.DeclRef;
import srt.ast.Expr;
import srt.ast.IntLiteral;
import srt.ast.Invariant;
import srt.ast.InvariantList;
import srt.ast.Program;
import srt.ast.WhileStmt;
import srt.ast.visitor.impl.DefaultVisitor;

// Visitor class implementation that generates loop candidate invariants
public class CandidatesGenVisitor extends DefaultVisitor {

	private List<WhileStmt> whileNodes; // all while nodes in given program
	private List<AssignStmt> assignmentNodes; // all assignment nodes in given program
	private Set<String> variableNames; // all variables present in given program
	private List<AssertStmt> assertNodes; // all assertions present in original program

	public CandidatesGenVisitor() {
		super(true);
	}

	//Generates invariants for given program
	@Override
	public Object visit(Program program) {
		//Information collection using CollectConstraintsVisitor
		CollectConstraintsVisitor ccv = new CollectConstraintsVisitor();
		ccv.visit(program);

		whileNodes = ccv.whileNodes;
		assignmentNodes = ccv.transitionNodes;
		variableNames = ccv.variableNames;
		assertNodes = ccv.propertyNodes;

		Set<Invariant> generatedInvariants; //invariants for loop
		List<WhileStmt> newWhileStmts = new ArrayList<WhileStmt>(); //container for changed loop statements
		Expr expr; 

		for (WhileStmt whileStmt : whileNodes) {
			//adding all already existing invariants to container
			generatedInvariants = new HashSet<Invariant>(whileStmt
					.getInvariantList().getInvariants());

			//adding while condidition candidate invariant
			generatedInvariants.add(new Invariant(true, whileStmt
					.getCondition()));

			//adding all invariants connected to assertions
			for (AssertStmt assertStmt : assertNodes) {
				generatedInvariants.add(new Invariant(true, assertStmt.getCondition()));
				
				try{
					expr = (BinaryExpr) assertStmt.getCondition();
					generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.EQUAL, ((BinaryExpr) expr).getLhs(), ((BinaryExpr) expr).getRhs())));
					generatedInvariants.add(new Invariant(true, new BinaryExpr(
							BinaryExpr.GEQ, ((BinaryExpr) expr).getLhs(), ((BinaryExpr) expr).getRhs())));
					generatedInvariants.add(new Invariant(true, new BinaryExpr(
							BinaryExpr.LEQ, ((BinaryExpr) expr).getLhs(), ((BinaryExpr) expr).getRhs())));
				} catch (ClassCastException e){
					
				}
				
			}


			for (AssignStmt assignStmt : assignmentNodes) {
				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.EQUAL, assignStmt.getLhs(), assignStmt
								.getRhs())));

				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.GEQ, assignStmt.getLhs(), assignStmt
								.getRhs())));

				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.LEQ, assignStmt.getLhs(), assignStmt
								.getRhs())));
			}
			
			//adding all invariants connected to assertions
			for (String s : variableNames) {
				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.GEQ, new DeclRef(s), new IntLiteral(0))));

				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.LEQ, new DeclRef(s), new IntLiteral(0))));

				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.GT, new DeclRef(s), new IntLiteral(0))));

				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.LT, new DeclRef(s), new IntLiteral(0))));

				generatedInvariants.add(new Invariant(true, new BinaryExpr(
						BinaryExpr.EQUAL, new DeclRef(s), new IntLiteral(0))));
				
				generatedInvariants.add(new Invariant(true, new BinaryExpr(BinaryExpr.EQUAL, new BinaryExpr(
						BinaryExpr.MOD, new DeclRef(s), new IntLiteral(2)), new IntLiteral(0))));
				
				generatedInvariants.add(new Invariant(true, new BinaryExpr(BinaryExpr.EQUAL, new BinaryExpr(
						BinaryExpr.MOD, new DeclRef(s), new IntLiteral(2)), new IntLiteral(1))));

				//adding all invariants connected to variables
				for (String s1 : variableNames) {
					if (s != s1) {
						generatedInvariants.add(new Invariant(true,
								new BinaryExpr(BinaryExpr.GEQ, new DeclRef(s),
										new DeclRef(s1))));

						generatedInvariants.add(new Invariant(true,
								new BinaryExpr(BinaryExpr.LEQ, new DeclRef(s),
										new DeclRef(s1))));

						generatedInvariants.add(new Invariant(true,
								new BinaryExpr(BinaryExpr.EQUAL,
										new DeclRef(s), new DeclRef(s1))));
					}
				}
			}

			//adding changed while statement to while statements container
			newWhileStmts.add(new WhileStmt(whileStmt.getCondition(), whileStmt
					.getBound(), new InvariantList(new ArrayList<Invariant>(
					generatedInvariants)), whileStmt.getBody()));
		}

		//update AST tree with new while nodes
		program = (Program) new ReplacementVisitor(newWhileStmts)
				.visit(program);

		return program;

	}

}
